from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Convolution2D, Activation, MaxPooling2D, Flatten, Dense
from tensorflow.keras.optimizers import Adam, SGD, Adagrad
from matplotlib import pyplot as plt

import mnist_reader
import os


# 获取优化器
def get_optimizer(lr=0.0001, opt='adam'):
    adam = Adam(lr=lr)  # Adam优化器
    sgd = SGD(lr=lr)  # SGD 优化器
    adagrad = Adagrad(learning_rate=lr)  # SGD 优化器
    optimizer_dict = {"adam": adam, "sgd": sgd, "adagrad": adagrad}
    opti = optimizer_dict[opt]
    return opti


def get_model1(opt=get_optimizer()):
    # 初始化模型
    model1 = Sequential()
    # 1st
    model1.add(Convolution2D(
        filters=32,  # 32个卷积核
        kernel_size=(5, 5),  # 卷积核设置(5,5)
        padding='same',  # 步长
        input_shape=(28, 28, 1)  # 输入图片的维度
    ))
    model1.add(Activation('relu'))
    model1.add(MaxPooling2D(
        pool_size=(2, 2),  # 设置2*2的核
        strides=(2, 2),  # 跳过2*2执行
        padding='same',  # 步长
    ))  # 输出(28,28,32)

    # 2nd
    model1.add(Convolution2D(
        filters=64,  # 64个卷积核
        kernel_size=(5, 5),  # 卷积核设置(5,5)
        padding='same',  # 步长
        input_shape=(28, 28, 1)  # 输入图片的维度
    ))
    model1.add(Activation('relu'))
    model1.add(MaxPooling2D(
        pool_size=(2, 2),  # 设置2*2的核
        strides=(2, 2),  # 跳过2*2执行
        padding='same',  # 步长
    ))

    # 1st Dense
    model1.add(Flatten())
    model1.add(Dense(64))
    model1.add(Activation('relu'))

    # 2nd Dense 最终输出10个选择
    model1.add(Dense(10))  # tabel
    model1.add(Activation('softmax'))

    # compile model1
    model1.compile(
        optimizer=opt,  # 指定优化器
        loss='categorical_crossentropy',
        metrics=['accuracy'],
    )

    return model1


# 比上面减少了一个卷积层一个池化层
def get_model2(opt=get_optimizer()):
    # 初始化模型
    model2 = Sequential()

    # 1st 卷积层
    model2.add(Convolution2D(
        filters=32,  # 32个卷积核
        kernel_size=(5, 5),  # 卷积核大小
        padding='same',  # 填充方式
        input_shape=(28, 28, 1)  # 输入图片的维度
    ))
    model2.add(Activation('relu'))
    model2.add(MaxPooling2D(
        pool_size=(2, 2),  # 池化核大小
        strides=(2, 2),  # 步长
        padding='same'  # 填充方式
    ))

    # 展平层
    model2.add(Flatten())

    # 1st 全连接层
    model2.add(Dense(64))
    model2.add(Activation('relu'))

    # 2nd 全连接层（输出层）
    model2.add(Dense(10))  # 输出层
    model2.add(Activation('softmax'))

    # 编译模型
    model2.compile(
        optimizer=opt,  # 指定优化器
        loss='categorical_crossentropy',  # 损失函数
        metrics=['accuracy']  # 评估指标
    )

    return model2


if __name__ == '__main__':
    # 先运行gpu_info看自己有没有cuda环境，注释掉这一行则用CPU跑
    # 配置GPU加速，根据gpu_info的可用GPU自己修改，比如现在是GPU0
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    # 配置参数
    nb_class = 10  # 代表10个类别
    nb_epoch = 10  # 迭代10次epoch, 一个epoch是将所有样本都过一遍，即1个epoch会过所有batch

    # 题目要求
    batchsize = 64  # 批次 根据自己的电脑配置设置
    learning_rate = 0.01  # 学习率 越小越慢 越大越容易过拟合
    optimizer = get_optimizer(lr=learning_rate, opt='adam')  # 优化器，可选项{"adam": adam, "sgd": sgd, "adagrad": adagrad}
    # 模型选择
    # model = get_model1(opt=optimizer)
    model = get_model2(opt=optimizer)

    # 1.读取数据集
    x_train, y_train = mnist_reader.load_mnist('data', kind='train')
    x_test, y_test = mnist_reader.load_mnist('data', kind='t10k')
    # print(len(x_train), len(y_train), len(x_test), len(y_test))

    # 2.修改数据格式
    # 标签转换为One-hot格式 [0,0,0,0,0,1,0,0,0,0] = 5
    y_train = np_utils.to_categorical(y_train, nb_class)
    y_test = np_utils.to_categorical(y_test, nb_class)

    # 调整图片形状
    x_train = x_train.reshape(-1, 28, 28, 1)  # -1代表未知，系统自己找，也可以自己设置60000，因为图片只有黑白通道所以设置1 RGB图片设置3
    x_test = x_test.reshape(-1, 28, 28, 1)

    # 3.训练模型
    history = model.fit(
        x=x_train,
        y=y_train,
        epochs=nb_epoch,
        batch_size=batchsize,
        validation_split=0.2,  # 20%训练集数据作为验证数据
        verbose=1
    )
    print()  # 换行，分开训练和测试
    # 4.评估模型
    scores = loss, acc = model.evaluate(
        x=x_test,
        y=y_test,
        batch_size=batchsize,
        verbose=1
    )
    # 打印测试结果
    print('test loss:', scores[0], ' accuracy:', scores[1])

    # 5.训练过程可视化
    # accuracy的历史
    plt.plot(history.history['accuracy'])
    plt.plot(history.history['val_accuracy'])
    plt.title('model accuracy')
    plt.ylabel('accuracy')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()

    # loss的历史
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()
